`timescale 1ns / 1ps

module ub_tb;

    // ========================================================================
    // Clock and Reset
    // ========================================================================
    logic clk;
    logic rst_n;

    // Clock generation: 10ns period (100MHz)
    initial begin
        clk = 0;
        forever #5 clk = ~clk;
    end

    // Reset generation
    initial begin
        rst_n = 0;
        repeat (10) @(posedge clk);  // Wait for 10 clocks before releasing reset
        rst_n = 1;
        $display("[%0t] Reset released", $time);
    end

    // ========================================================================
    // Interface Instantiations
    // ========================================================================

    // Read request interfaces
    ub_rd_req_if u_ldma_rd_req_if (
        .clk  (clk),
        .rst_n(rst_n)
    );

    ub_rd_req_if u_aru_rd_req_if (
        .clk  (clk),
        .rst_n(rst_n)
    );

    // Read data interfaces
    aru_ub_dat_if u_aru_rd_dat_if (
        .clk  (clk),
        .rst_n(rst_n)
    );

    ldma_ub_dat_if u_ldma_rd_dat_if (
        .clk  (clk),
        .rst_n(rst_n)
    );

    // Write request interfaces
    aru_ub_wr_req_if u_aru_wr_req_if (
        .clk  (clk),
        .rst_n(rst_n)
    );

    gdma_ub_wr_req_if u_gdma_wr_req_if (
        .clk  (clk),
        .rst_n(rst_n)
    );

    // ========================================================================
    // DUT Instantiation
    // ========================================================================
    ub u_dut (
        .clk             (clk),
        .rst_n           (rst_n),
        .u_ldma_rd_req_if(u_ldma_rd_req_if.in),
        .u_aru_rd_req_if (u_aru_rd_req_if.in),
        .u_aru_rd_dat_if (u_aru_rd_dat_if.out),
        .u_ldma_rd_dat_if(u_ldma_rd_dat_if.out),
        .u_aru_wr_req_if (u_aru_wr_req_if.in),
        .u_gdma_wr_req_if(u_gdma_wr_req_if.in)
    );

    // // ========================================================================
    // // Memory Models (Simplified BFMs)
    // // ========================================================================

    // // LMB Memory Model
    // // Parameters (following mxu_psb.sv style)
    // localparam LMB_BANK_WIDTH = `M0 * `K0_Byte;              // 16 * 16 = 256 bytes
    // localparam LMB_BANK_DEPTH = `LMB_SIZE / LMB_BANK_WIDTH;  // 65536 / 256 = 256

    // // LMB Storage - organized as M0 rows of pack_16B_t per entry
    // pack_16B_t lmb_mem [LMB_BANK_DEPTH-1:0][`M0-1:0];

    // always_ff @(posedge clk or negedge rst_n) begin
    //     if (!rst_n) begin
    //         lmb_rd_dat_if.vld <= 1'b0;
    //     end else begin
    //         lmb_rd_dat_if.rdy <= 1'b1;  // Always ready
    //         if (lmb_rd_req_if.vld && lmb_rd_req_if.rdy) begin
    //             lmb_rd_dat_if.vld <= 1'b1;
    //             lmb_rd_dat_if.dat <= lmb_mem[lmb_rd_req_if.lmb_addr];
    //             $display("[%0t] LMB Read: addr=0x%h", $time, lmb_rd_req_if.lmb_addr);
    //         end else begin
    //             lmb_rd_dat_if.vld <= 1'b0;
    //         end
    //     end
    // end
    // assign lmb_rd_req_if.rdy = 1'b1;

    // // RMB Memory Model
    // // Parameters (following mxu_psb.sv style)
    // localparam RMB_BANK_WIDTH = `N0 * `K0_Byte;              // 8 * 16 = 128 bytes
    // localparam RMB_BANK_DEPTH = `RMB_SIZE / RMB_BANK_WIDTH;  // 65536 / 128 = 512

    // // RMB Storage - organized as N0 columns of pack_16B_t per entry
    // pack_16B_t rmb_mem [RMB_BANK_DEPTH-1:0][`N0-1:0];

    // always_ff @(posedge clk or negedge rst_n) begin
    //     if (!rst_n) begin
    //         rmb_rd_dat_if.vld <= 1'b0;
    //     end else begin
    //         rmb_rd_dat_if.rdy <= 1'b1;
    //         if (rmb_rd_req_if.vld && rmb_rd_req_if.rdy) begin
    //             rmb_rd_dat_if.vld <= 1'b1;
    //             rmb_rd_dat_if.dat <= rmb_mem[rmb_rd_req_if.rmb_addr];
    //             $display("[%0t] RMB Read: addr=0x%h", $time, rmb_rd_req_if.rmb_addr);
    //         end else begin
    //             rmb_rd_dat_if.vld <= 1'b0;
    //         end
    //     end
    // end
    // assign rmb_rd_req_if.rdy = 1'b1;

    // // PMB Memory Model
    // // Parameters (following mxu_psb.sv style)
    // localparam PMB_BANK_WIDTH = `N0 * `ACC_Byte;             // 8 * 4 = 32 bytes
    // localparam PMB_BANK_DEPTH = `PMB_SIZE / PMB_BANK_WIDTH;  // 16384 / 32 = 512

    // // PMB Storage - organized as N0 fp32 elements per entry
    // fp32_t pmb_mem [PMB_BANK_DEPTH-1:0][`N0-1:0];
    // always_ff @(posedge clk or negedge rst_n) begin
    //     if (!rst_n) begin
    //         pmb_rd_dat_if.vld <= 1'b0;
    //     end else begin
    //         pmb_rd_dat_if.rdy <= 1'b1;
    //         if (pmb_rd_req_if.vld && pmb_rd_req_if.rdy) begin
    //             pmb_rd_dat_if.vld <= 1'b1;
    //             pmb_rd_dat_if.dat <= pmb_mem[pmb_rd_req_if.pmb_addr];
    //             $display("[%0t] PMB Read: addr=0x%h", $time, pmb_rd_req_if.pmb_addr);
    //         end else begin
    //             pmb_rd_dat_if.vld <= 1'b0;
    //         end
    //     end
    // end
    // assign pmb_rd_req_if.rdy = 1'b1;

    // // ARU PSB Read Model (simplified - no actual reads in this test)
    // assign aru_psb_rd_req_if.vld = 1'b0;
    // assign aru_psb_rd_req_if.psb_addr = 24'h0;

    // // ========================================================================
    // // Test Stimulus
    // // ========================================================================

    // // Task: Send MXU instruction
    // task send_mxu_instr(
    //     input logic        psum_en,
    //     input logic        bias_en,
    //     input logic [15:0] lmb_addr,
    //     input logic [15:0] rmb_addr,
    //     input logic [23:0] pmb_addr,
    //     input logic [23:0] psb_addr,
    //     input logic [15:0] m1,
    //     input logic [15:0] n1,
    //     input logic [15:0] k1
    // );
    //     @(posedge clk);
    //     lpu_mxu_instr_if.vld <= 1'b1;
    //     lpu_mxu_instr_if.pld.psum_en    <= psum_en;
    //     lpu_mxu_instr_if.pld.bias_en    <= bias_en;
    //     lpu_mxu_instr_if.pld.lmb_addr   <= lmb_addr;
    //     lpu_mxu_instr_if.pld.rmb_addr   <= rmb_addr;
    //     lpu_mxu_instr_if.pld.pmb_addr   <= pmb_addr;
    //     lpu_mxu_instr_if.pld.psb_addr   <= psb_addr;
    //     lpu_mxu_instr_if.pld.slice_m    <= m1;
    //     lpu_mxu_instr_if.pld.slice_n    <= n1;
    //     lpu_mxu_instr_if.pld.slice_k1   <= k1;

    //     // Wait for ready
    //     wait(lpu_mxu_instr_if.rdy);
    //     @(posedge clk);
    //     lpu_mxu_instr_if.vld <= 1'b0;

    //     $display("[%0t] Sent instruction: M1=%0d, N1=%0d, K1=%0d, psum_en=%b, bias_en=%b",
    //              $time, m1, n1, k1, psum_en, bias_en);
    // endtask

    // // Task: Initialize memory with test data
    // task init_memory();
    //     integer i, m, n, k;

    //     $display("[%0t] Initializing memory with test data...", $time);

    //     // Initialize LMB with simple pattern
    //     // LMB: [DEPTH][M0] where each element is pack_16B_t containing K0 bf16 values
    //     for (i = 0; i < LMB_BANK_DEPTH; i++) begin
    //         for (m = 0; m < `M0; m++) begin
    //             // Use bf16 format: 8 bf16 values per pack_16B_t (K0=8)
    //             for (k = 0; k < `K0; k++) begin
    //                 lmb_mem[i][m].bf16[k] = 16'h3f80;  // bf16: 1.0
    //             end
    //         end
    //     end

    //     // Initialize RMB with simple pattern
    //     // RMB: [DEPTH][N0] where each element is pack_16B_t containing K0 bf16 values
    //     for (i = 0; i < RMB_BANK_DEPTH; i++) begin
    //         for (n = 0; n < `N0; n++) begin
    //             // Use bf16 format: 8 bf16 values per pack_16B_t (K0=8)
    //             for (k = 0; k < `K0; k++) begin
    //                 rmb_mem[i][n].bf16[k] = 16'h3f80;  // bf16: 1.0
    //             end
    //         end
    //     end

    //     // Initialize PMB with bias values
    //     // PMB: [DEPTH][N0] where each element is fp32_t
    //     for (i = 0; i < PMB_BANK_DEPTH; i++) begin
    //         for (n = 0; n < `N0; n++) begin
    //             pmb_mem[i][n] = 32'h3f800000;  // fp32: 1.0
    //         end
    //     end

    //     $display("[%0t] Memory initialization complete", $time);
    // endtask

    // // Task: Monitor done signal
    // task wait_for_done();
    //     $display("[%0t] Waiting for computation to complete...", $time);
    //     wait(done_if.in.vld);
    //     @(posedge clk);
    //     $display("[%0t] Computation DONE!", $time);
    // endtask

    // // ========================================================================
    // // Main Test Sequence
    // // ========================================================================
    // initial begin
    //     // Initialize
    //     lpu_mxu_instr_if.vld <= 1'b0;

    //     // Wait for reset
    //     wait(rst_n);
    //     repeat(5) @(posedge clk);

    //     // Initialize test data
    //     init_memory();

    //     // ====================================================================
    //     // Test Case 1: Simple Matrix Multiplication (no psum, no bias)
    //     // ====================================================================
    //     $display("\n========================================");
    //     $display("Test Case 1: A * B (M1=2, N1=2, K1=2)");
    //     $display("========================================");
    //     send_mxu_instr(
    //         .psum_en  (1'b0),
    //         .bias_en  (1'b0),
    //         .lmb_addr (16'h0000),
    //         .rmb_addr (16'h0000),
    //         .pmb_addr (24'h000000),
    //         .psb_addr (24'h000000),
    //         .m1       (16'd2),
    //         .n1       (16'd2),
    //         .k1       (16'd2)
    //     );
    //     wait_for_done();

    //     repeat(10) @(posedge clk);

    //     // ====================================================================
    //     // Test Case 2: Matrix Multiplication with Bias
    //     // ====================================================================
    //     $display("\n========================================");
    //     $display("Test Case 2: A * B + bias (M1=1, N1=1, K1=1)");
    //     $display("========================================");
    //     send_mxu_instr(
    //         .psum_en  (1'b0),
    //         .bias_en  (1'b1),
    //         .lmb_addr (16'h0100),
    //         .rmb_addr (16'h0100),
    //         .pmb_addr (24'h000100),
    //         .psb_addr (24'h000100),
    //         .m1       (16'd1),
    //         .n1       (16'd1),
    //         .k1       (16'd1)
    //     );
    //     wait_for_done();

    //     repeat(10) @(posedge clk);

    //     // ====================================================================
    //     // Test Case 3: Matrix Multiplication with Partial Sum
    //     // ====================================================================
    //     $display("\n========================================");
    //     $display("Test Case 3: A * B + psum (M1=1, N1=1, K1=2)");
    //     $display("========================================");
    //     send_mxu_instr(
    //         .psum_en  (1'b1),
    //         .bias_en  (1'b0),
    //         .lmb_addr (16'h0200),
    //         .rmb_addr (16'h0200),
    //         .pmb_addr (24'h000200),
    //         .psb_addr (24'h000200),
    //         .m1       (16'd1),
    //         .n1       (16'd1),
    //         .k1       (16'd2)
    //     );
    //     wait_for_done();

    //     repeat(20) @(posedge clk);

    //     // ====================================================================
    //     // End of Test
    //     // ====================================================================
    //     $display("\n========================================");
    //     $display("All tests completed successfully!");
    //     $display("========================================");
    //     repeat(50) @(posedge clk);
    //     $finish;
    // end

    // // ========================================================================
    // // Timeout Watchdog
    // // ========================================================================
    // initial begin
    //     #1000000;  // 1ms timeout
    //     $display("\n[ERROR] Timeout! Test did not complete in time.");
    //     $finish;
    // end

    // // ========================================================================
    // // Waveform Dump
    // // ========================================================================
    // initial begin
    //     $dumpfile("mxu_tb.vcd");
    //     $dumpvars(0, mxu_tb);
    // end

endmodule
